{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer 实战学习"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 包引入"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import copy\n",
    "from importlib import import_module\n",
    "import torch.nn as nn\n",
    "from sklearn import metrics\n",
    "import time\n",
    "from tensorboardX import SummaryWriter\n",
    "import os\n",
    "import pickle as pkl\n",
    "from tqdm import tqdm\n",
    "import time\n",
    "from datetime import timedelta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 配置项"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "code_folding": []
   },
   "outputs": [],
   "source": [
    "dataset = 'THUCNews'  # 数据集\n",
    "model_name = 'Transformer'\n",
    "embedding = 'embedding_SougouNews.npz'    # 搜狗新闻:embedding_SougouNews.npz, 腾讯:embedding_Tencent.npz, 随机初始化:random\n",
    "word = False     # True for word, False for char\n",
    "class Config(object):\n",
    "\n",
    "    \"\"\"配置参数\"\"\"\n",
    "    def __init__(self, dataset, embedding):\n",
    "        self.model_name = 'Transformer'\n",
    "        self.train_path = dataset + '/data/train.txt'                                # 训练集\n",
    "        self.dev_path = dataset + '/data/dev.txt'                                    # 验证集\n",
    "        self.test_path = dataset + '/data/test.txt'                                  # 测试集\n",
    "        self.class_list = [x.strip() for x in open(\n",
    "            dataset + '/data/class.txt', encoding='utf-8').readlines()]              # 类别名单\n",
    "        self.vocab_path = dataset + '/data/vocab.pkl'                                # 词表\n",
    "        self.save_path = dataset + '/saved_dict/' + self.model_name + '.ckpt'        # 模型训练结果\n",
    "        self.log_path = dataset + '/log/' + self.model_name\n",
    "        self.embedding_pretrained = torch.tensor(\n",
    "            np.load(dataset + '/data/' + embedding)[\"embeddings\"].astype('float32'))\\\n",
    "            if embedding != 'random' else None                                       # 预训练词向量\n",
    "        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')   # 设备\n",
    "\n",
    "        self.dropout = 0.5                                              # 随机失活\n",
    "        self.require_improvement = 2000                                 # 若超过1000batch效果还没提升，则提前结束训练\n",
    "        self.num_classes = len(self.class_list)                         # 类别数\n",
    "        self.n_vocab = 0                                                # 词表大小，在运行时赋值\n",
    "        self.num_epochs = 20                                            # epoch数\n",
    "        self.batch_size = 128                                           # mini-batch大小\n",
    "        self.pad_size = 32                                              # 每句话处理成的长度(短填长切)\n",
    "        self.learning_rate = 5e-4                                       # 学习率\n",
    "        self.embed = self.embedding_pretrained.size(1)\\\n",
    "            if self.embedding_pretrained is not None else 300           # 字向量维度\n",
    "        self.dim_model = 300\n",
    "        self.hidden = 1024\n",
    "        self.last_hidden = 512\n",
    "        self.num_head = 5\n",
    "        self.num_encoder = 2 \n",
    "        \n",
    "        \n",
    "config = Config(dataset, embedding)\n",
    "np.random.seed(1)\n",
    "torch.manual_seed(1)\n",
    "torch.cuda.manual_seed_all(1)\n",
    "torch.backends.cudnn.deterministic = True  # 保证每次结果一样"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 数据加载"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "heading_collapsed": true
   },
   "source": [
    "### 函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "code_folding": [
     3,
     48,
     54,
     133,
     136,
     143,
     154
    ],
    "hidden": true
   },
   "outputs": [],
   "source": [
    "MAX_VOCAB_SIZE = 10000  # 词表长度限制\n",
    "UNK, PAD = '<UNK>', '<PAD>'  # 未知字，padding符号\n",
    "# 功能：构建词典\n",
    "def build_vocab(file_path, tokenizer, max_size, min_freq):\n",
    "    '''\n",
    "        功能：构建词典\n",
    "        input:\n",
    "            file_path      String      数据集目录\n",
    "            tokenizer      Function    数据集 分词 or 分字 处理函数\n",
    "            max_size       int         词典最大长度   \n",
    "            min_freq       int         词频\n",
    "        output:\n",
    "            vocab_dic      Dict        所构建的词典\n",
    "    '''\n",
    "    vocab_dic = {}\n",
    "    with open(file_path, 'r', encoding='UTF-8') as f:\n",
    "        for line in tqdm(f):\n",
    "            lin = line.strip()\n",
    "            if not lin:\n",
    "                continue\n",
    "            content = lin.split('\\t')[0]\n",
    "            for word in tokenizer(content):\n",
    "                vocab_dic[word] = vocab_dic.get(word, 0) + 1\n",
    "        vocab_list = sorted([_ for _ in vocab_dic.items() if _[1] >= min_freq], key=lambda x: x[1], reverse=True)[:max_size]\n",
    "        vocab_dic = {word_count[0]: idx for idx, word_count in enumerate(vocab_list)}\n",
    "        vocab_dic.update({UNK: len(vocab_dic), PAD: len(vocab_dic) + 1})\n",
    "    return vocab_dic\n",
    "# 功能：处理数据集\n",
    "def build_dataset(config, ues_word):\n",
    "    '''\n",
    "        功能：处理数据集\n",
    "        input:\n",
    "            config        Object      配置项        \n",
    "            ues_word      Boolean     True for word, False for char （如果是 词，需要事先 分词）\n",
    "        output:\n",
    "            vocab         Dict        所构建的词典\n",
    "            train         List        [(([...], 0),seq_len), (([...], 1),seq_len), ...]\n",
    "            dev           List        [(([...], 0),seq_len), (([...], 1),seq_len), ...]\n",
    "            test          List        [(([...], 0),seq_len), (([...], 1),seq_len), ...] \n",
    "    '''\n",
    "    # 词或字符 预处理操作\n",
    "    if ues_word:\n",
    "        tokenizer = lambda x: x.split(' ')  # 以空格隔开，word-level\n",
    "    else:\n",
    "        tokenizer = lambda x: [y for y in x]  # char-level\n",
    "    # 加载词典\n",
    "    if os.path.exists(config.vocab_path):\n",
    "        vocab = pkl.load(open(config.vocab_path, 'rb'))\n",
    "    else:\n",
    "        vocab = build_vocab(config.train_path, tokenizer=tokenizer, max_size=MAX_VOCAB_SIZE, min_freq=1)\n",
    "        pkl.dump(vocab, open(config.vocab_path, 'wb'))\n",
    "    print(f\"Vocab size: {len(vocab)}\")\n",
    "    \n",
    "    # 加载 数据集\n",
    "    def load_dataset(path, pad_size=32):\n",
    "        '''\n",
    "            功能：加载 数据集 \n",
    "            input:\n",
    "                path         String      数据集地址 \n",
    "                pad_size     int         填充方式 or 句子 max_len\n",
    "            output:\n",
    "                contents:    List        [(([...], 0),seq_len), (([...], 1),seq_len), ...]\n",
    "        '''\n",
    "        contents = []\n",
    "        with open(path, 'r', encoding='UTF-8') as f:\n",
    "            for line in tqdm(f):\n",
    "                lin = line.strip()\n",
    "                if not lin:\n",
    "                    continue\n",
    "                content, label = lin.split('\\t')\n",
    "                words_line = []\n",
    "                token = tokenizer(content)\n",
    "                seq_len = len(token)\n",
    "                if pad_size:\n",
    "                    if len(token) < pad_size:\n",
    "                        token.extend([PAD] * (pad_size - len(token)))\n",
    "                    else:\n",
    "                        token = token[:pad_size]\n",
    "                        seq_len = pad_size\n",
    "                # word to id\n",
    "                for word in token:\n",
    "                    words_line.append(vocab.get(word, vocab.get(UNK)))\n",
    "                contents.append((words_line, int(label), seq_len))\n",
    "        return contents  # [(([...], 0),seq_len), (([...], 1),seq_len), ...]\n",
    "    train = load_dataset(config.train_path, config.pad_size)\n",
    "    dev = load_dataset(config.dev_path, config.pad_size)\n",
    "    test = load_dataset(config.test_path, config.pad_size)\n",
    "    return vocab, train, dev, test\n",
    "# 功能：构建 数据 iterator Class\n",
    "class DatasetIterater(object):\n",
    "    def __init__(self, batches, batch_size, device):\n",
    "        '''\n",
    "            功能：构建 数据 iterator Class\n",
    "            input:\n",
    "                batches         List        [(([...], 0),seq_len), (([...], 1),seq_len), ...]\n",
    "                batch_size      int         batch 大小  \n",
    "                device          Object      GPU or CPU\n",
    "            output:\n",
    "                \n",
    "        '''\n",
    "        self.batch_size = batch_size\n",
    "        self.batches = batches\n",
    "        self.n_batches = len(batches) // batch_size\n",
    "        self.residue = False  # 记录batch数量是否为整数\n",
    "        if len(batches) % self.n_batches != 0:\n",
    "            self.residue = True\n",
    "        self.index = 0\n",
    "        self.device = device\n",
    "    # 功能：转 tensor \n",
    "    def _to_tensor(self, datas):\n",
    "        x = torch.LongTensor([_[0] for _ in datas]).to(self.device)\n",
    "        y = torch.LongTensor([_[1] for _ in datas]).to(self.device)\n",
    "\n",
    "        # pad前的长度(超过pad_size的设为pad_size)\n",
    "        seq_len = torch.LongTensor([_[2] for _ in datas]).to(self.device)\n",
    "        return (x, seq_len), y\n",
    "\n",
    "    def __next__(self):\n",
    "        if self.residue and self.index == self.n_batches:\n",
    "            batches = self.batches[self.index * self.batch_size: len(self.batches)]\n",
    "            self.index += 1\n",
    "            batches = self._to_tensor(batches)\n",
    "            return batches\n",
    "\n",
    "        elif self.index >= self.n_batches:\n",
    "            self.index = 0\n",
    "            raise StopIteration\n",
    "        else:\n",
    "            batches = self.batches[self.index * self.batch_size: (self.index + 1) * self.batch_size]\n",
    "            self.index += 1\n",
    "            batches = self._to_tensor(batches)\n",
    "            return batches\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __len__(self):\n",
    "        if self.residue:\n",
    "            return self.n_batches + 1\n",
    "        else:\n",
    "            return self.n_batches\n",
    "# 功能：构建 数据 iterator\n",
    "def build_iterator(dataset, config):\n",
    "    '''\n",
    "        功能：构建 数据 iterator\n",
    "        input:\n",
    "            dataset         List        [(([...], 0),seq_len), (([...], 1),seq_len), ...]\n",
    "            config          Object      配置项   \n",
    "        output:\n",
    "            iter            iterator    迭代器\n",
    "    '''\n",
    "    iter = DatasetIterater(dataset, config.batch_size, config.device)\n",
    "    return iter\n",
    "\n",
    "def get_time_dif(start_time):\n",
    "    \"\"\"获取已使用时间\"\"\"\n",
    "    end_time = time.time()\n",
    "    time_dif = end_time - start_time\n",
    "    return timedelta(seconds=int(round(time_dif)))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 数据集 处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading data...\n",
      "Vocab size: 4762\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "180000it [00:03, 55226.60it/s]\n",
      "10000it [00:00, 61126.40it/s]\n",
      "10000it [00:00, 58963.84it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Time usage: 0:00:04\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "print(\"Loading data...\")\n",
    "vocab, train_data, dev_data, test_data = build_dataset(config, word)\n",
    "train_iter = build_iterator(train_data, config)\n",
    "dev_iter = build_iterator(dev_data, config)\n",
    "test_iter = build_iterator(test_data, config)\n",
    "time_dif = get_time_dif(start_time)\n",
    "print(\"Time usage:\", time_dif)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型构建"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 模型构建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model(nn.Module):\n",
    "    def __init__(self, config):\n",
    "        super(Model, self).__init__()\n",
    "        if config.embedding_pretrained is not None:\n",
    "            self.embedding = nn.Embedding.from_pretrained(config.embedding_pretrained, freeze=False)\n",
    "        else:\n",
    "            self.embedding = nn.Embedding(config.n_vocab, config.embed, padding_idx=config.n_vocab - 1)\n",
    "\n",
    "        self.postion_embedding = Positional_Encoding(config.embed, config.pad_size, config.dropout, config.device)\n",
    "        self.encoder = Encoder(config.dim_model, config.num_head, config.hidden, config.dropout)\n",
    "        self.encoders = nn.ModuleList([\n",
    "            copy.deepcopy(self.encoder)\n",
    "            # Encoder(config.dim_model, config.num_head, config.hidden, config.dropout)\n",
    "            for _ in range(config.num_encoder)])\n",
    "\n",
    "        self.fc1 = nn.Linear(config.pad_size * config.dim_model, config.num_classes)\n",
    "        # self.fc2 = nn.Linear(config.last_hidden, config.num_classes)\n",
    "        # self.fc1 = nn.Linear(config.dim_model, config.num_classes)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.embedding(x[0])\n",
    "        out = self.postion_embedding(out)\n",
    "        for encoder in self.encoders:\n",
    "            out = encoder(out)\n",
    "        out = out.view(out.size(0), -1)\n",
    "        # out = torch.mean(out, 1)\n",
    "        out = self.fc1(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Encoder(nn.Module):\n",
    "    def __init__(self, dim_model, num_head, hidden, dropout):\n",
    "        super(Encoder, self).__init__()\n",
    "        self.attention = Multi_Head_Attention(dim_model, num_head, dropout)\n",
    "        self.feed_forward = Position_wise_Feed_Forward(dim_model, hidden, dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.attention(x)\n",
    "        out = self.feed_forward(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Positional_Encoding(nn.Module):\n",
    "    def __init__(self, embed, pad_size, dropout, device):\n",
    "        super(Positional_Encoding, self).__init__()\n",
    "        self.device = device\n",
    "        self.pe = torch.tensor([[pos / (10000.0 ** (i // 2 * 2.0 / embed)) for i in range(embed)] for pos in range(pad_size)])\n",
    "        self.pe[:, 0::2] = np.sin(self.pe[:, 0::2])\n",
    "        self.pe[:, 1::2] = np.cos(self.pe[:, 1::2])\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = x + nn.Parameter(self.pe, requires_grad=False).to(self.device)\n",
    "        out = self.dropout(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Scaled_Dot_Product_Attention(nn.Module):\n",
    "    '''Scaled Dot-Product Attention '''\n",
    "    def __init__(self):\n",
    "        super(Scaled_Dot_Product_Attention, self).__init__()\n",
    "\n",
    "    def forward(self, Q, K, V, scale=None):\n",
    "        '''\n",
    "        Args:\n",
    "            Q: [batch_size, len_Q, dim_Q]\n",
    "            K: [batch_size, len_K, dim_K]\n",
    "            V: [batch_size, len_V, dim_V]\n",
    "            scale: 缩放因子 论文为根号dim_K\n",
    "        Return:\n",
    "            self-attention后的张量，以及attention张量\n",
    "        '''\n",
    "        attention = torch.matmul(Q, K.permute(0, 2, 1))\n",
    "        if scale:\n",
    "            attention = attention * scale\n",
    "        # if mask:  # TODO change this\n",
    "        #     attention = attention.masked_fill_(mask == 0, -1e9)\n",
    "        attention = F.softmax(attention, dim=-1)\n",
    "        context = torch.matmul(attention, V)\n",
    "        return context\n",
    "\n",
    "\n",
    "class Multi_Head_Attention(nn.Module):\n",
    "    def __init__(self, dim_model, num_head, dropout=0.0):\n",
    "        super(Multi_Head_Attention, self).__init__()\n",
    "        self.num_head = num_head\n",
    "        assert dim_model % num_head == 0\n",
    "        self.dim_head = dim_model // self.num_head\n",
    "        self.fc_Q = nn.Linear(dim_model, num_head * self.dim_head)\n",
    "        self.fc_K = nn.Linear(dim_model, num_head * self.dim_head)\n",
    "        self.fc_V = nn.Linear(dim_model, num_head * self.dim_head)\n",
    "        self.attention = Scaled_Dot_Product_Attention()\n",
    "        self.fc = nn.Linear(num_head * self.dim_head, dim_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.layer_norm = nn.LayerNorm(dim_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        batch_size = x.size(0)\n",
    "        Q = self.fc_Q(x)\n",
    "        K = self.fc_K(x)\n",
    "        V = self.fc_V(x)\n",
    "        Q = Q.view(batch_size * self.num_head, -1, self.dim_head)\n",
    "        K = K.view(batch_size * self.num_head, -1, self.dim_head)\n",
    "        V = V.view(batch_size * self.num_head, -1, self.dim_head)\n",
    "        # if mask:  # TODO\n",
    "        #     mask = mask.repeat(self.num_head, 1, 1)  # TODO change this\n",
    "        scale = K.size(-1) ** -0.5  # 缩放因子\n",
    "        context = self.attention(Q, K, V, scale)\n",
    "\n",
    "        context = context.view(batch_size, -1, self.dim_head * self.num_head)\n",
    "        out = self.fc(context)\n",
    "        out = self.dropout(out)\n",
    "        out = out + x  # 残差连接\n",
    "        out = self.layer_norm(out)\n",
    "        return out\n",
    "\n",
    "\n",
    "class Position_wise_Feed_Forward(nn.Module):\n",
    "    def __init__(self, dim_model, hidden, dropout=0.0):\n",
    "        super(Position_wise_Feed_Forward, self).__init__()\n",
    "        self.fc1 = nn.Linear(dim_model, hidden)\n",
    "        self.fc2 = nn.Linear(hidden, dim_model)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        self.layer_norm = nn.LayerNorm(dim_model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        out = self.fc1(x)\n",
    "        out = F.relu(out)\n",
    "        out = self.fc2(out)\n",
    "        out = self.dropout(out)\n",
    "        out = out + x  # 残差连接\n",
    "        out = self.layer_norm(out)\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 模型初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<bound method Module.parameters of Model(\n",
      "  (embedding): Embedding(4762, 300)\n",
      "  (postion_embedding): Positional_Encoding(\n",
      "    (dropout): Dropout(p=0.5)\n",
      "  )\n",
      "  (encoder): Encoder(\n",
      "    (attention): Multi_Head_Attention(\n",
      "      (fc_Q): Linear(in_features=300, out_features=300, bias=True)\n",
      "      (fc_K): Linear(in_features=300, out_features=300, bias=True)\n",
      "      (fc_V): Linear(in_features=300, out_features=300, bias=True)\n",
      "      (attention): Scaled_Dot_Product_Attention()\n",
      "      (fc): Linear(in_features=300, out_features=300, bias=True)\n",
      "      (dropout): Dropout(p=0.5)\n",
      "      (layer_norm): LayerNorm(torch.Size([300]), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "    (feed_forward): Position_wise_Feed_Forward(\n",
      "      (fc1): Linear(in_features=300, out_features=1024, bias=True)\n",
      "      (fc2): Linear(in_features=1024, out_features=300, bias=True)\n",
      "      (dropout): Dropout(p=0.5)\n",
      "      (layer_norm): LayerNorm(torch.Size([300]), eps=1e-05, elementwise_affine=True)\n",
      "    )\n",
      "  )\n",
      "  (encoders): ModuleList(\n",
      "    (0): Encoder(\n",
      "      (attention): Multi_Head_Attention(\n",
      "        (fc_Q): Linear(in_features=300, out_features=300, bias=True)\n",
      "        (fc_K): Linear(in_features=300, out_features=300, bias=True)\n",
      "        (fc_V): Linear(in_features=300, out_features=300, bias=True)\n",
      "        (attention): Scaled_Dot_Product_Attention()\n",
      "        (fc): Linear(in_features=300, out_features=300, bias=True)\n",
      "        (dropout): Dropout(p=0.5)\n",
      "        (layer_norm): LayerNorm(torch.Size([300]), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "      (feed_forward): Position_wise_Feed_Forward(\n",
      "        (fc1): Linear(in_features=300, out_features=1024, bias=True)\n",
      "        (fc2): Linear(in_features=1024, out_features=300, bias=True)\n",
      "        (dropout): Dropout(p=0.5)\n",
      "        (layer_norm): LayerNorm(torch.Size([300]), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "    )\n",
      "    (1): Encoder(\n",
      "      (attention): Multi_Head_Attention(\n",
      "        (fc_Q): Linear(in_features=300, out_features=300, bias=True)\n",
      "        (fc_K): Linear(in_features=300, out_features=300, bias=True)\n",
      "        (fc_V): Linear(in_features=300, out_features=300, bias=True)\n",
      "        (attention): Scaled_Dot_Product_Attention()\n",
      "        (fc): Linear(in_features=300, out_features=300, bias=True)\n",
      "        (dropout): Dropout(p=0.5)\n",
      "        (layer_norm): LayerNorm(torch.Size([300]), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "      (feed_forward): Position_wise_Feed_Forward(\n",
      "        (fc1): Linear(in_features=300, out_features=1024, bias=True)\n",
      "        (fc2): Linear(in_features=1024, out_features=300, bias=True)\n",
      "        (dropout): Dropout(p=0.5)\n",
      "        (layer_norm): LayerNorm(torch.Size([300]), eps=1e-05, elementwise_affine=True)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (fc1): Linear(in_features=9600, out_features=10, bias=True)\n",
      ")>\n"
     ]
    }
   ],
   "source": [
    "config.n_vocab = len(vocab)\n",
    "model = Model(config).to(config.device)\n",
    "if model_name != 'Transformer':\n",
    "    init_network(model)\n",
    "print(model.parameters)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型训练"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 函数编写"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "code_folding": [
     1,
     17,
     71,
     87
    ]
   },
   "outputs": [],
   "source": [
    "# 权重初始化，默认xavier\n",
    "def init_network(model, method='xavier', exclude='embedding', seed=123):\n",
    "    for name, w in model.named_parameters():\n",
    "        if exclude not in name:\n",
    "            if 'weight' in name:\n",
    "                if method == 'xavier':\n",
    "                    nn.init.xavier_normal_(w)\n",
    "                elif method == 'kaiming':\n",
    "                    nn.init.kaiming_normal_(w)\n",
    "                else:\n",
    "                    nn.init.normal_(w)\n",
    "            elif 'bias' in name:\n",
    "                nn.init.constant_(w, 0)\n",
    "            else:\n",
    "                pass\n",
    "\n",
    "def train(config, model, train_iter, dev_iter, test_iter):\n",
    "    start_time = time.time()\n",
    "    model.train()\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)\n",
    "\n",
    "    # 学习率指数衰减，每次epoch：学习率 = gamma * 学习率\n",
    "    # scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.9)\n",
    "    total_batch = 0  # 记录进行到多少batch\n",
    "    dev_best_loss = float('inf')\n",
    "    last_improve = 0  # 记录上次验证集loss下降的batch数\n",
    "    flag = False  # 记录是否很久没有效果提升\n",
    "    writer = SummaryWriter(log_dir=config.log_path + '/' + time.strftime('%m-%d_%H.%M', time.localtime()))\n",
    "    for epoch in range(config.num_epochs):\n",
    "        print('Epoch [{}/{}]'.format(epoch + 1, config.num_epochs))\n",
    "        # scheduler.step() # 学习率衰减\n",
    "        for i, (trains, labels) in enumerate(train_iter):\n",
    "            outputs = model(trains)\n",
    "            model.zero_grad()\n",
    "            loss = F.cross_entropy(outputs, labels)\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "            if total_batch % 100 == 0:\n",
    "                # 每多少轮输出在训练集和验证集上的效果\n",
    "                true = labels.data.cpu()\n",
    "                predic = torch.max(outputs.data, 1)[1].cpu()\n",
    "                train_acc = metrics.accuracy_score(true, predic)\n",
    "                dev_acc, dev_loss = evaluate(config, model, dev_iter)\n",
    "                if dev_loss < dev_best_loss:\n",
    "                    dev_best_loss = dev_loss\n",
    "                    torch.save(model.state_dict(), config.save_path)\n",
    "                    improve = '*'\n",
    "                    last_improve = total_batch\n",
    "                else:\n",
    "                    improve = ''\n",
    "                time_dif = get_time_dif(start_time)\n",
    "                msg = 'Iter: {0:>6},  Train Loss: {1:>5.2},  Train Acc: {2:>6.2%},  Val Loss: {3:>5.2},  Val Acc: {4:>6.2%},  Time: {5} {6}'\n",
    "                print(msg.format(total_batch, loss.item(), train_acc, dev_loss, dev_acc, time_dif, improve))\n",
    "                writer.add_scalar(\"loss/train\", loss.item(), total_batch)\n",
    "                writer.add_scalar(\"loss/dev\", dev_loss, total_batch)\n",
    "                writer.add_scalar(\"acc/train\", train_acc, total_batch)\n",
    "                writer.add_scalar(\"acc/dev\", dev_acc, total_batch)\n",
    "                model.train()\n",
    "            total_batch += 1\n",
    "            if total_batch - last_improve > config.require_improvement:\n",
    "                # 验证集loss超过1000batch没下降，结束训练\n",
    "                print(\"No optimization for a long time, auto-stopping...\")\n",
    "                flag = True\n",
    "                break\n",
    "        if flag:\n",
    "            break\n",
    "    writer.close()\n",
    "    test(config, model, test_iter)\n",
    "\n",
    "def test(config, model, test_iter):\n",
    "    # test\n",
    "    model.load_state_dict(torch.load(config.save_path))\n",
    "    model.eval()\n",
    "    start_time = time.time()\n",
    "    test_acc, test_loss, test_report, test_confusion = evaluate(config, model, test_iter, test=True)\n",
    "    msg = 'Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}'\n",
    "    print(msg.format(test_loss, test_acc))\n",
    "    print(\"Precision, Recall and F1-Score...\")\n",
    "    print(test_report)\n",
    "    print(\"Confusion Matrix...\")\n",
    "    print(test_confusion)\n",
    "    time_dif = get_time_dif(start_time)\n",
    "    print(\"Time usage:\", time_dif)\n",
    "\n",
    "def evaluate(config, model, data_iter, test=False):\n",
    "    model.eval()\n",
    "    loss_total = 0\n",
    "    predict_all = np.array([], dtype=int)\n",
    "    labels_all = np.array([], dtype=int)\n",
    "    with torch.no_grad():\n",
    "        for texts, labels in data_iter:\n",
    "            outputs = model(texts)\n",
    "            loss = F.cross_entropy(outputs, labels)\n",
    "            loss_total += loss\n",
    "            labels = labels.data.cpu().numpy()\n",
    "            predic = torch.max(outputs.data, 1)[1].cpu().numpy()\n",
    "            labels_all = np.append(labels_all, labels)\n",
    "            predict_all = np.append(predict_all, predic)\n",
    "\n",
    "    acc = metrics.accuracy_score(labels_all, predict_all)\n",
    "    if test:\n",
    "        report = metrics.classification_report(labels_all, predict_all, target_names=config.class_list, digits=4)\n",
    "        confusion = metrics.confusion_matrix(labels_all, predict_all)\n",
    "        return acc, loss_total / len(data_iter), report, confusion\n",
    "    return acc, loss_total / len(data_iter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 函数调用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch [1/20]\n",
      "Iter:      0,  Train Loss:   2.4,  Train Acc: 10.94%,  Val Loss:   4.6,  Val Acc: 10.01%,  Time: 0:00:20 *\n"
     ]
    }
   ],
   "source": [
    "train(config, model, train_iter, dev_iter, test_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.3"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
